package com.riteny.oauth.core.tokenstore;

import com.riteny.oauth.core.userauth.CustomUser;
import com.riteny.oauth.core.userauth.CustomUserAuthenticationService;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.store.InMemoryTokenStore;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author Riteny
 * 2021/5/7  17:39
 */
public class CustomInMemoryTokenStore extends InMemoryTokenStore implements CustomTokenStore {

    private final ConcurrentHashMap<String, Map<String, Object>> userInfoStore = new ConcurrentHashMap();

    private CustomUserAuthenticationService customUserAuthenticationService;

    public CustomInMemoryTokenStore(CustomUserAuthenticationService customUserAuthenticationService) {
        this.customUserAuthenticationService = customUserAuthenticationService;
    }

    @Override
    public void storeAccessToken(OAuth2AccessToken token, OAuth2Authentication authentication) {

        super.storeAccessToken(token, authentication);

        CustomUser map = (CustomUser) authentication.getUserAuthentication().getPrincipal();
        String userType = map.getUserType();
        String username = map.getUsername();

        Map<String, Object> userInfo = new HashMap<>();
        userInfo.put("userDetails", customUserAuthenticationService.searchUserDetailInfo(username, userType));
        userInfoStore.put(token.getValue(), userInfo);
    }

    @Override
    public Map<String, Object> readUserInfo(OAuth2AccessToken token, OAuth2Authentication authentication) {
        return this.userInfoStore.get(token);
    }

    @Override
    public void removeAccessToken(OAuth2AccessToken accessToken) {

        super.removeAccessToken(accessToken);

        userInfoStore.remove(accessToken.getValue());
    }

    @Override
    public void removeAccessToken(String tokenValue) {

        super.removeAccessToken(tokenValue);

        userInfoStore.remove(tokenValue);
    }
}
